from torch.utils.data import Dataset, DataLoader
import os
import pandas as pd
import numpy as np
import torch


class iris_dataloader(Dataset):
    def __init__(self, data_pass):
        self.data_pass = data_pass

        assert os.path.exists(self.data_pass), "data_pass does not exist"
        df = pd.read_csv(self.data_pass, names=[0, 1, 2, 3, 4])

        d = {"Iris-setosa": 0, "Iris-versicolor": 1, "Iris-virginica": 2}
        df[4] = df[4].map(d)

        data = df.iloc[:,0:4]
        label = df.iloc[:,4:]
        means = data.mean(axis=0)  # 按列计算均值
        stds = data.std(axis=0)  # 按列计算标准差
        data = (data - means) / stds
        # data = ((data - np.mean(data))/np.std(data))
        # data = (data - np.mean(data)) / np.std(data)

        self.data = torch.from_numpy(np.array(data,dtype='float32'))
        self.label = torch.from_numpy(np.array(label,dtype='int64'))

        self.data_num = len(label)
        print(f'当前数据集大小：{self.data_num}')

    def __len__(self):
        return self.data_num

    def __getitem__(self, index):
        self.data = list(self.data)
        self.label = list(self.label)
        return self.data[index], self.label[index]
